{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/force.html).  I recommend having both windows open, side-by-side!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    Box,\n",
    "    ConstantVectorSource,\n",
    "    ContactVisualizer,\n",
    "    ContactVisualizerParams,\n",
    "    DiagramBuilder,\n",
    "    DiscreteContactSolver,\n",
    "    FixedOffsetFrame,\n",
    "    InverseDynamicsController,\n",
    "    JointSliders,\n",
    "    LeafSystem,\n",
    "    LogVectorOutput,\n",
    "    MathematicalProgram,\n",
    "    MeshcatVisualizer,\n",
    "    MeshcatVisualizerParams,\n",
    "    MultibodyPlant,\n",
    "    MultibodyPositionToGeometryPose,\n",
    "    Multiplexer,\n",
    "    OsqpSolver,\n",
    "    Parser,\n",
    "    PiecewisePolynomial,\n",
    "    PlanarJoint,\n",
    "    PrismaticJoint,\n",
    "    RevoluteJoint,\n",
    "    RigidTransform,\n",
    "    RotationMatrix,\n",
    "    SceneGraph,\n",
    "    Simulator,\n",
    "    SpatialInertia,\n",
    "    Sphere,\n",
    "    StartMeshcat,\n",
    "    TrajectorySource,\n",
    "    UnitInertia,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.meshcat_utils import MeshcatSliders\n",
    "from manipulation.scenarios import AddShape, SetColor, SetTransparency\n",
    "from manipulation.utils import ConfigureParser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gOFIUWF2KBPN"
   },
   "outputs": [],
   "source": [
    "def AddPlanarBinAndCracker(plant):\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    bin = parser.AddModelsFromUrl(\"package://manipulation/planar_bin.sdf\")[0]\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"bin_base\", bin),\n",
    "        RigidTransform(\n",
    "            RotationMatrix.MakeZRotation(np.pi / 2.0), [0, 0, -0.015]\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    planar_joint_frame = plant.AddFrame(\n",
    "        FixedOffsetFrame(\n",
    "            \"planar_joint_frame\",\n",
    "            plant.world_frame(),\n",
    "            RigidTransform(RotationMatrix.MakeXRotation(np.pi / 2)),\n",
    "        )\n",
    "    )\n",
    "    box = parser.AddModelsFromUrl(\n",
    "        \"package://drake/manipulation/models/ycb/sdf/003_cracker_box.sdf\"\n",
    "    )[0]\n",
    "    box_frame = plant.AddFrame(\n",
    "        FixedOffsetFrame(\n",
    "            \"box_frame\",\n",
    "            plant.GetFrameByName(\"base_link_cracker\", box),\n",
    "            RigidTransform(RotationMatrix.MakeXRotation(np.pi / 2)),\n",
    "        )\n",
    "    )\n",
    "    box_joint = plant.AddJoint(\n",
    "        PlanarJoint(\"box_joint\", planar_joint_frame, box_frame)\n",
    "    )\n",
    "    box_joint.set_default_translation([0, 0.033400])\n",
    "    return box\n",
    "\n",
    "\n",
    "def AddPlanarBinAndSimpleBox(\n",
    "    plant, mass=1.0, mu=1.0, width=0.2, depth=0.05, height=0.3\n",
    "):\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    bin = parser.AddModelsFromUrl(\"package://manipulation/planar_bin.sdf\")[0]\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"bin_base\", bin),\n",
    "        RigidTransform(\n",
    "            RotationMatrix.MakeZRotation(np.pi / 2.0), [0, 0, -0.015]\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    planar_joint_frame = plant.AddFrame(\n",
    "        FixedOffsetFrame(\n",
    "            \"planar_joint_frame\",\n",
    "            plant.world_frame(),\n",
    "            RigidTransform(RotationMatrix.MakeXRotation(np.pi / 2)),\n",
    "        )\n",
    "    )\n",
    "\n",
    "    # TODO(russt): make this a *random* box?\n",
    "    # TODO(russt): move random box to a shared py file.\n",
    "    box_instance = AddShape(plant, Box(width, depth, height), \"box\", mass, mu)\n",
    "    box_joint = plant.AddJoint(\n",
    "        PlanarJoint(\n",
    "            \"box_joint\",\n",
    "            planar_joint_frame,\n",
    "            plant.GetFrameByName(\"box\", box_instance),\n",
    "        )\n",
    "    )\n",
    "    box_joint.set_default_translation([0, height / 2.0])\n",
    "    return box_instance\n",
    "\n",
    "\n",
    "def AddPlanarGripper(plant):\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    gripper = parser.AddModelsFromUrl(\n",
    "        \"package://manipulation/schunk_wsg_50_welded_fingers.sdf\"\n",
    "    )[0]\n",
    "    gripper_body = plant.GetBodyByName(\"body\", gripper)\n",
    "\n",
    "    # Add a planar joint the old fashioned way (so that I can have three actuators):\n",
    "    gripper_false_body1 = plant.AddRigidBody(\n",
    "        \"false_body1\",\n",
    "        gripper,\n",
    "        SpatialInertia(0, [0, 0, 0], UnitInertia(0, 0, 0)),\n",
    "    )\n",
    "    gripper_false_body2 = plant.AddRigidBody(\n",
    "        \"false_body2\",\n",
    "        gripper,\n",
    "        SpatialInertia(0, [0, 0, 0], UnitInertia(0, 0, 0)),\n",
    "    )\n",
    "    gripper_x = plant.AddJoint(\n",
    "        PrismaticJoint(\n",
    "            \"gripper_x\",\n",
    "            plant.world_frame(),\n",
    "            plant.GetFrameByName(\"false_body1\"),\n",
    "            [1, 0, 0],\n",
    "            -0.3,\n",
    "            0.3,\n",
    "        )\n",
    "    )\n",
    "    plant.AddJointActuator(\"gripper_x\", gripper_x)\n",
    "    gripper_z = plant.AddJoint(\n",
    "        PrismaticJoint(\n",
    "            \"gripper_z\",\n",
    "            plant.GetFrameByName(\"false_body1\"),\n",
    "            plant.GetFrameByName(\"false_body2\"),\n",
    "            [0, 0, 1],\n",
    "            0.0,\n",
    "            0.5,\n",
    "        )\n",
    "    )\n",
    "    gripper_z.set_default_translation(0.3)\n",
    "    plant.AddJointActuator(\"gripper_z\", gripper_z)\n",
    "    gripper_frame = plant.AddFrame(\n",
    "        FixedOffsetFrame(\n",
    "            \"gripper_frame\",\n",
    "            plant.GetFrameByName(\"body\", gripper),\n",
    "            RigidTransform(RotationMatrix.MakeXRotation(np.pi / 2)),\n",
    "        )\n",
    "    )\n",
    "    gripper_theta = plant.AddJoint(\n",
    "        RevoluteJoint(\n",
    "            \"gripper_theta\",\n",
    "            plant.GetFrameByName(\"false_body2\"),\n",
    "            gripper_frame,\n",
    "            [0, -1, 0],\n",
    "            -np.pi,\n",
    "            np.pi,\n",
    "        )\n",
    "    )\n",
    "    plant.AddJointActuator(\"gripper_theta\", gripper_theta)\n",
    "\n",
    "    return gripper\n",
    "\n",
    "\n",
    "def MakePlanarGripperOnlyPlant():\n",
    "    plant = MultibodyPlant(time_step=0.005)\n",
    "    AddPlanarGripper(plant)\n",
    "    plant.Finalize()\n",
    "    return plant\n",
    "\n",
    "\n",
    "def AddPointFinger(plant):\n",
    "    finger = AddShape(\n",
    "        plant, Sphere(0.01), \"finger\", color=[0.9, 0.5, 0.5, 1.0]\n",
    "    )\n",
    "    false_body1 = plant.AddRigidBody(\n",
    "        \"false_body1\",\n",
    "        finger,\n",
    "        SpatialInertia(0, [0, 0, 0], UnitInertia(0, 0, 0)),\n",
    "    )\n",
    "    finger_x = plant.AddJoint(\n",
    "        PrismaticJoint(\n",
    "            \"finger_x\",\n",
    "            plant.world_frame(),\n",
    "            plant.GetFrameByName(\"false_body1\"),\n",
    "            [1, 0, 0],\n",
    "            -0.3,\n",
    "            0.3,\n",
    "        )\n",
    "    )\n",
    "    plant.AddJointActuator(\"finger_x\", finger_x)\n",
    "    finger_z = plant.AddJoint(\n",
    "        PrismaticJoint(\n",
    "            \"finger_z\",\n",
    "            plant.GetFrameByName(\"false_body1\"),\n",
    "            plant.GetFrameByName(\"finger\"),\n",
    "            [0, 0, 1],\n",
    "            0.0,\n",
    "            0.5,\n",
    "        )\n",
    "    )\n",
    "    finger_z.set_default_translation(0.25)\n",
    "    plant.AddJointActuator(\"finger_z\", finger_z)\n",
    "\n",
    "    return finger\n",
    "\n",
    "\n",
    "def MakeFingerOnlyPlant():\n",
    "    plant = MultibodyPlant(time_step=0.005)\n",
    "    AddPointFinger(plant)\n",
    "    plant.Finalize()\n",
    "    return plant\n",
    "\n",
    "\n",
    "class PointFingerForceControl(LeafSystem):\n",
    "    def __init__(self, plant):\n",
    "        LeafSystem.__init__(self)\n",
    "        self._plant = plant\n",
    "\n",
    "        self.DeclareVectorInputPort(\"desired_contact_force\", 2)\n",
    "        self.DeclareVectorOutputPort(\"finger_actuation\", 2, self.CalcOutput)\n",
    "\n",
    "    def CalcOutput(self, context, output):\n",
    "        finger_mass = 1\n",
    "        g = self._plant.gravity_field().gravity_vector()[[0, 2]]\n",
    "\n",
    "        desired_force = self.get_input_port(0).Eval(context)\n",
    "        output.SetFromVector(-finger_mass * g - desired_force)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Commanding a constant force"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def force_control_point_finger(\n",
    "    desired_horizontal_forces=[4.0], duration=1.0, draw=False\n",
    "):\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    box = AddPlanarBinAndCracker(plant)\n",
    "    finger = AddPointFinger(plant)\n",
    "    plant.GetJointByName(\"finger_x\").set_default_translation(0.15)\n",
    "    plant.GetJointByName(\"finger_z\").set_default_translation(0.05)\n",
    "    plant.Finalize()\n",
    "\n",
    "    if draw:\n",
    "        MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "        #        vis.vis['drake']['4'].delete()  # Only flush the gripper geometry.\n",
    "        #    vis.set_planar_viewpoint(xmin=-.5, xmax=.5, ymin=-0.1, ymax=0.6)\n",
    "        ContactVisualizer.AddToBuilder(\n",
    "            builder, plant, meshcat, ContactVisualizerParams(radius=0.005)\n",
    "        )\n",
    "    controller = builder.AddSystem(PointFingerForceControl(plant))\n",
    "    builder.Connect(\n",
    "        controller.get_output_port(), plant.get_actuation_input_port()\n",
    "    )\n",
    "    builder.ExportInput(controller.get_input_port(), \"desired_contact_force\")\n",
    "\n",
    "    logger = LogVectorOutput(plant.get_state_output_port(finger), builder)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    plt.figure()\n",
    "    lines = []\n",
    "    for f in desired_horizontal_forces:\n",
    "        simulator = Simulator(diagram)\n",
    "        context = simulator.get_mutable_context()\n",
    "        diagram.get_input_port().FixValue(context, [f, 0])\n",
    "\n",
    "        simulator.AdvanceTo(duration)\n",
    "        log = logger.FindMutableLog(context)\n",
    "        t = log.sample_times()\n",
    "        x = log.data()\n",
    "        lines.append(plt.plot(t, x[0, :])[0])\n",
    "\n",
    "    plt.legend(\n",
    "        lines,\n",
    "        [\n",
    "            (\"$f^{F_c}_{desired}$ = \" + str(f))\n",
    "            for f in desired_horizontal_forces\n",
    "        ],\n",
    "    )\n",
    "    plt.xlabel(\"t (sec)\")\n",
    "    plt.ylabel(\"x (m)\")\n",
    "\n",
    "\n",
    "force_control_point_finger([0.25, 0.5, 1.0, 2.0, 3.0, 4.0, 5.0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VBc7zDKlElKu"
   },
   "source": [
    "#  A force-based flip-up strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xJ7sDtzbEfYw",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def force_control_point_finger(teleop=True):\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.005)\n",
    "    plant.set_discrete_contact_solver(DiscreteContactSolver.kSap)\n",
    "    box = AddPlanarBinAndCracker(plant)\n",
    "    finger = AddPointFinger(plant)\n",
    "    plant.GetJointByName(\"finger_x\").set_default_translation(0.15)\n",
    "    plant.GetJointByName(\"finger_z\").set_default_translation(0.025)\n",
    "    plant.Finalize()\n",
    "\n",
    "    vis = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "    # vis.vis['drake']['4'].delete()  # Only flush the gripper geometry.\n",
    "    # vis.set_planar_viewpoint(xmin=-.5, xmax=.5, ymin=-0.1, ymax=0.6)\n",
    "    contact_vis = ContactVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        plant,\n",
    "        meshcat,\n",
    "        ContactVisualizerParams(radius=0.005, newtons_per_meter=40.0),\n",
    "    )\n",
    "\n",
    "    class RotateBoxOnTopCorner(LeafSystem):\n",
    "        def __init__(self, plant):\n",
    "            LeafSystem.__init__(self)\n",
    "            self._plant = plant\n",
    "\n",
    "            # One state variable for the integral error.\n",
    "            self.DeclareContinuousState(1)\n",
    "\n",
    "            self.DeclareVectorInputPort(\"desired_box_angle\", 1)\n",
    "            self.DeclareVectorInputPort(\"box_state\", 6)\n",
    "            self.DeclareVectorInputPort(\"finger_state\", 4)\n",
    "            self.DeclareVectorOutputPort(\"finger_force\", 2, self.CalcOutput)\n",
    "            self.solver = OsqpSolver()\n",
    "\n",
    "        def DoCalcTimeDerivatives(self, context, derivatives):\n",
    "            # Compute integral of the angle error.\n",
    "            theta_desired = self.get_input_port(0).Eval(context)[0]\n",
    "            theta = self.get_input_port(1).Eval(context)[2]\n",
    "            derivatives.get_mutable_vector().SetFromVector(\n",
    "                [theta_desired - theta]\n",
    "            )\n",
    "\n",
    "        def CalcOutput(self, context, output):\n",
    "            # TODO(russt): Clean this up to match the cleaned up derivation in\n",
    "            # the text.\n",
    "            box_mass = 0.411  # Let's revisit whether we need this!\n",
    "            mu_A = 1.0\n",
    "            mu_C = 1.0\n",
    "            w = 0.2  # Approximate box width\n",
    "            g = -self._plant.gravity_field().gravity_vector()[2]\n",
    "\n",
    "            theta_desired = self.get_input_port(0).Eval(context)[0]\n",
    "            box_state = self.get_input_port(1).Eval(context)\n",
    "            theta = box_state[2]\n",
    "            theta_dot = box_state[5]\n",
    "\n",
    "            R_WC = np.array(\n",
    "                [\n",
    "                    [-np.sin(theta), -np.cos(theta)],\n",
    "                    [np.cos(theta), -np.sin(theta)],\n",
    "                ]\n",
    "            )\n",
    "            R_CW = R_WC.T\n",
    "            finger_state = self.get_input_port(2).Eval(context)\n",
    "\n",
    "            prog = MathematicalProgram()\n",
    "            f_C_W = prog.NewContinuousVariables(\n",
    "                2, \"f_C_W\"\n",
    "            )  # force from finger to box\n",
    "            f_C_C = np.matmul(R_CW, f_C_W)\n",
    "\n",
    "            # PID control on f_CC[1]\n",
    "            kp = 20.0\n",
    "            ki = 0.5\n",
    "            kd = 2.0 * np.sqrt(kp)  # should really be kp*I on inside\n",
    "            integral_error = context.get_continuous_state_vector().GetAtIndex(\n",
    "                0\n",
    "            )\n",
    "            pid = (\n",
    "                kp * (theta_desired - theta)\n",
    "                + ki * integral_error\n",
    "                - kd * theta_dot\n",
    "            )\n",
    "            prog.AddQuadraticCost((w * f_C_C[0] - pid) ** 2, True)\n",
    "\n",
    "            prog.AddConstraint(f_C_W[0] <= mu_A * (box_mass * g - f_C_W[1]))\n",
    "            prog.AddConstraint(-f_C_W[0] <= mu_A * (box_mass * g - f_C_W[1]))\n",
    "            prog.AddConstraint(\n",
    "                f_C_C[1] >= 0.1\n",
    "            )  # To ensure I make and stay in contact.\n",
    "            prog.AddConstraint(f_C_C[0] <= mu_C * f_C_C[1])\n",
    "            prog.AddConstraint(-f_C_C[0] <= mu_C * f_C_C[1])\n",
    "\n",
    "            result = self.solver.Solve(prog)\n",
    "\n",
    "            if not result.is_success():\n",
    "                print(result.GetInfeasibleConstraintNames(prog))\n",
    "                output.SetFromVector([0, 0])\n",
    "                return\n",
    "\n",
    "            f_C_W = result.GetSolution(f_C_W)\n",
    "            output.SetFromVector(-f_C_W)\n",
    "\n",
    "    controller = builder.AddSystem(PointFingerForceControl(plant))\n",
    "    builder.Connect(\n",
    "        controller.get_output_port(), plant.get_actuation_input_port()\n",
    "    )\n",
    "\n",
    "    rotate_box = builder.AddSystem(RotateBoxOnTopCorner(plant))\n",
    "    builder.Connect(rotate_box.get_output_port(), controller.get_input_port())\n",
    "    builder.Connect(\n",
    "        plant.get_state_output_port(box), rotate_box.get_input_port(1)\n",
    "    )\n",
    "    builder.Connect(\n",
    "        plant.get_state_output_port(finger), rotate_box.get_input_port(2)\n",
    "    )\n",
    "\n",
    "    if teleop:\n",
    "        meshcat.AddSlider(\n",
    "            \"desired box angle\",\n",
    "            value=0,\n",
    "            min=0,\n",
    "            max=np.pi / 2,\n",
    "            step=0.05,\n",
    "            increment_keycode=\"ArrowLeft\",\n",
    "            decrement_keycode=\"ArrowRight\",\n",
    "        )\n",
    "        teleop = builder.AddSystem(\n",
    "            MeshcatSliders(meshcat, [[\"desired box angle\"]])\n",
    "        )\n",
    "        builder.Connect(teleop.get_output_port(), rotate_box.get_input_port(0))\n",
    "    else:\n",
    "        theta_trajectory = PiecewisePolynomial.FirstOrderHold(\n",
    "            [0, 2, 20], [[0, 0, np.pi / 2.0]]\n",
    "        )\n",
    "        trajectory_source = builder.AddSystem(\n",
    "            TrajectorySource(theta_trajectory)\n",
    "        )\n",
    "        builder.Connect(\n",
    "            trajectory_source.get_output_port(), rotate_box.get_input_port(0)\n",
    "        )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "\n",
    "    if not running_as_notebook:\n",
    "        simulator.AdvanceTo(0.01)\n",
    "    elif teleop:\n",
    "        meshcat.AddButton(\"Stop Simulation\", \"Escape\")\n",
    "        while meshcat.GetButtonClicks(\"Stop Simulation\") < 1:\n",
    "            simulator.AdvanceTo(simulator.get_context().get_time() + 2.0)\n",
    "        meshcat.DeleteButton(\"Stop Simulation\")\n",
    "        meshcat.DeleteAddedControls()\n",
    "    else:\n",
    "        simulator.set_target_realtime_rate(0.5)  # slow motion!\n",
    "        vis.StartRecording()\n",
    "        simulator.AdvanceTo(theta_trajectory.end_time())\n",
    "        vis.StopRecording()\n",
    "        contact_vis.Delete()  # TODO: make my recording include the contact forces!\n",
    "        vis.PublishRecording()\n",
    "\n",
    "\n",
    "force_control_point_finger(teleop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7q0A14bAilIX"
   },
   "source": [
    "# Teleop with Stiffness Control\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JX1Oq7gfWsBD"
   },
   "outputs": [],
   "source": [
    "def stiffness_box_flipup(point_finger=True, teleop=False):\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    # plant.set_discrete_contact_solver(DiscreteContactSolver.kSap)\n",
    "\n",
    "    AddPlanarBinAndSimpleBox(plant)\n",
    "    # AddPlanarBinAndCracker(plant)\n",
    "    SetTransparency(scene_graph, alpha=0.5, source_id=plant.get_source_id())\n",
    "\n",
    "    gripper = (\n",
    "        AddPointFinger(plant) if point_finger else AddPlanarGripper(plant)\n",
    "    )\n",
    "    plant.Finalize()\n",
    "\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "    meshcat.Set2dRenderMode(xmin=-0.5, xmax=0.5, ymin=-0.1, ymax=0.6)\n",
    "    ContactVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        plant,\n",
    "        meshcat,\n",
    "        ContactVisualizerParams(radius=0.005, newtons_per_meter=40.0),\n",
    "    )\n",
    "\n",
    "    # Use the controller plant to visualize the set point geometry.\n",
    "    controller_scene_graph = builder.AddSystem(SceneGraph())\n",
    "    controller_plant = MultibodyPlant(time_step=0.005)\n",
    "    controller_plant.RegisterAsSourceForSceneGraph(controller_scene_graph)\n",
    "    AddPointFinger(controller_plant) if point_finger else AddPlanarGripper(\n",
    "        controller_plant\n",
    "    )\n",
    "    controller_plant.Finalize()\n",
    "    SetColor(\n",
    "        controller_scene_graph,\n",
    "        color=[1.0, 165.0 / 255, 0.0, 1.0],\n",
    "        source_id=controller_plant.get_source_id(),\n",
    "    )\n",
    "    controller_vis = MeshcatVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        controller_scene_graph,\n",
    "        meshcat,\n",
    "        MeshcatVisualizerParams(prefix=\"controller\"),\n",
    "    )\n",
    "    controller_vis.set_name(\"controller meshcat\")\n",
    "\n",
    "    # TODO(russt): Implement the proper JointStiffnessController in drake and\n",
    "    # use it here. (For a point finger with unit mass, the\n",
    "    # InverseDynamicsController is identical)\n",
    "    N = controller_plant.num_positions()\n",
    "    kp = [100] * N\n",
    "    ki = [1] * N\n",
    "    kd = [2 * np.sqrt(kp[0])] * N\n",
    "    controller = builder.AddSystem(\n",
    "        InverseDynamicsController(controller_plant, kp, ki, kd, False)\n",
    "    )\n",
    "    builder.Connect(\n",
    "        plant.get_state_output_port(gripper),\n",
    "        controller.get_input_port_estimated_state(),\n",
    "    )\n",
    "\n",
    "    positions_to_state = builder.AddSystem(Multiplexer([N, N]))\n",
    "    zeros = builder.AddSystem(ConstantVectorSource([0] * N))\n",
    "    builder.Connect(\n",
    "        zeros.get_output_port(0), positions_to_state.get_input_port(1)\n",
    "    )\n",
    "    builder.Connect(\n",
    "        positions_to_state.get_output_port(),\n",
    "        controller.get_input_port_desired_state(),\n",
    "    )\n",
    "    positions_to_poses = builder.AddSystem(\n",
    "        MultibodyPositionToGeometryPose(controller_plant)\n",
    "    )\n",
    "    builder.Connect(\n",
    "        positions_to_poses.get_output_port(),\n",
    "        controller_scene_graph.get_source_pose_port(\n",
    "            controller_plant.get_source_id()\n",
    "        ),\n",
    "    )\n",
    "    builder.Connect(\n",
    "        controller.get_output_port_control(), plant.get_actuation_input_port()\n",
    "    )\n",
    "\n",
    "    if teleop:\n",
    "        position_sliders = builder.AddSystem(\n",
    "            JointSliders(\n",
    "                meshcat,\n",
    "                controller_plant,\n",
    "                step=0.01,\n",
    "                decrement_keycodes=[\"ArrowLeft\", \"ArrowDown\"],\n",
    "                increment_keycodes=[\"ArrowRight\", \"ArrowUp\"],\n",
    "            )\n",
    "        )\n",
    "        builder.Connect(\n",
    "            position_sliders.get_output_port(),\n",
    "            positions_to_state.get_input_port(0),\n",
    "        )\n",
    "        builder.Connect(\n",
    "            position_sliders.get_output_port(),\n",
    "            positions_to_poses.get_input_port(),\n",
    "        )\n",
    "    else:\n",
    "        traj = PiecewisePolynomial.FirstOrderHold(\n",
    "            [0, 10.0, 20.0],\n",
    "            np.array([[0.15, 0.025], [-0.25, 0.025], [-0.25, 0.125]]).T,\n",
    "        )\n",
    "        traj_system = builder.AddSystem(TrajectorySource(traj))\n",
    "        builder.Connect(\n",
    "            traj_system.get_output_port(), positions_to_state.get_input_port(0)\n",
    "        )\n",
    "        builder.Connect(\n",
    "            traj_system.get_output_port(), positions_to_poses.get_input_port()\n",
    "        )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "\n",
    "    if not running_as_notebook:\n",
    "        simulator.AdvanceTo(0.01)\n",
    "    elif teleop:\n",
    "        simulator.set_target_realtime_rate(1.0)\n",
    "        meshcat.AddButton(\"Stop Simulation\", \"Escape\")\n",
    "        while meshcat.GetButtonClicks(\"Stop Simulation\") < 1:\n",
    "            simulator.AdvanceTo(simulator.get_context().get_time() + 2.0)\n",
    "        meshcat.DeleteButton(\"Stop Simulation\")\n",
    "        meshcat.DeleteAddedControls()\n",
    "    else:\n",
    "        simulator.set_target_realtime_rate(1.0)\n",
    "        simulator.AdvanceTo(traj.end_time() + 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A stiffness-control-based flip-up strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stiffness_box_flipup(point_finger=True, teleop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stiffness_box_flipup(point_finger=True, teleop=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "g3lNMaloEfY3"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Force Control.ipynb",
   "provenance": []
  },
  "interpreter": {
   "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
